import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from scipy.stats import pointbiserialr
import numpy as np


def plot_distribution_by_col1(df, col1, col2, binning=True):
    if binning:
        # 将col1转换为数值类型，如果不能转换则为NaN
        df[col1] = pd.to_numeric(df[col1], errors='coerce')

        # 将col1等分成20份
        df['col1_bins'] = pd.cut(df[col1], bins=20)

        # 获取col1_bins的所有唯一值
        col1_bins = df['col1_bins'].unique()

        # 创建一个图形
        fig, axes = plt.subplots(5, 4, figsize=(20, 15), sharex=True, sharey=True)
        axes = axes.flatten()

        for i, col1_bin in enumerate(col1_bins):
            if i >= len(axes):
                break

            # 过滤出当前col1_bin的所有行
            filtered_df = df[df['col1_bins'] == col1_bin]

            # 统计每个col2_bin的出现次数
            col2_distribution = filtered_df[col2].value_counts().sort_index()

            # 绘制col2的直方图
            col2_distribution.plot(kind='bar', ax=axes[i], edgecolor='black')

            # 在直方图上添加文本标签
            for idx, count in enumerate(col2_distribution):
                if count > 0:  # 只在有数据的bin上添加标签
                    bin_label = col2_distribution.index[idx]
                    axes[i].text(idx, count, f'{bin_label}\n{int(count)}', ha='center', va='bottom', fontsize=8,
                                 color='black', rotation=90)

            axes[i].set_title(f'{col1}: {col1_bin}')
            axes[i].set_xlabel(col2)
            axes[i].set_ylabel('Frequency')

    else:
        # 获取col1的所有唯一值
        col1_unique_values = df[col1].unique()

        # 创建一个图形
        fig, axes = plt.subplots(5, 4, figsize=(20, 15), sharex=True, sharey=True)
        axes = axes.flatten()

        for i, col1_value in enumerate(col1_unique_values):
            if i >= len(axes):
                break

            # 过滤出当前col1_value的所有行
            filtered_df = df[df[col1] == col1_value]

            # 统计每个col2_bin的出现次数
            col2_distribution = filtered_df[col2].value_counts().sort_index()

            # 绘制col2的直方图
            col2_distribution.plot(kind='bar', ax=axes[i], edgecolor='black')

            # 在直方图上添加文本标签
            for idx, count in enumerate(col2_distribution):
                if count > 0:  # 只在有数据的bin上添加标签
                    bin_label = col2_distribution.index[idx]
                    axes[i].text(idx, count, f'{bin_label}\n{int(count)}', ha='center', va='bottom', fontsize=8,
                                 color='black', rotation=90)

            axes[i].set_title(f'{col1}: {col1_value}')
            axes[i].set_xlabel(col2)
            axes[i].set_ylabel('Frequency')

    plt.tight_layout()
    plt.show()
def plot_correlation(result_df):

    # 将日期列转换为datetime格式
    result_df['trading_date'] = pd.to_datetime(result_df['trading_date'])

    # 设置日期列为索引
    result_df.set_index('trading_date', inplace=True)

    # 获取所有的code
    codes = result_df['code1'].unique()

    # 绘制折线图
    plt.figure(figsize=(30, 6))

    for label, df_group in result_df.groupby(['code1', 'code2']):
        plt.plot( df_group['corr'], label='-'.join(label))

    # for code in codes:
    #     code_df = result_df[result_df['code1'] == code]
    #     plt.plot(code_df.index, code_df['corr'], label=f'Correlation ({code})')
        # plt.plot(code_df.index, code_df['pearson_corr'], label=f'Pearson Correlation ({code})')
        # plt.plot(code_df.index, code_df['spearman_corr'], label=f'Spearman Correlation ({code})')
    #
    # for code in codes:
    #     code_df = result_df[result_df['code1'] == code]
    #
    #     # 使用滑动平均平滑数据
    #     code_df_smoothed = code_df.rolling(window=3, min_periods=1).mean()
    #
    #     plt.plot(code_df_smoothed.index, code_df_smoothed['pearson_corr'], label=f'Pearson Correlation ({code})')
    #     plt.plot(code_df_smoothed.index, code_df_smoothed['spearman_corr'], label=f'Spearman Correlation ({code})')


    plt.title('Correlation Over Time')
    plt.xlabel('Date')
    plt.ylabel('Correlation')
    plt.legend()
    plt.grid(True)
    plt.xticks(rotation=45)

    # 设置x轴刻度为每5天
    plt.gca().xaxis.set_major_locator(mdates.DayLocator(interval=10))

    plt.tight_layout()
    plt.show()


def plot_histograms(df, column_names):
    """
    绘制指定列的直方图，每个直方图将数据分为20份。

    参数:
    df: pd.DataFrame
        数据框
    column_names: list of str
        列名列表
    """
    num_columns = len(column_names)

    # 设置图形的尺寸
    fig, axes = plt.subplots(num_columns, 1, figsize=(10, 5 * num_columns))

    if num_columns == 1:
        axes = [axes]  # 保证axes是一个可迭代对象，即使只有一个子图

    for i, column in enumerate(column_names):
        ax = axes[i]
        df[column].plot(kind='hist', bins=40, ax=ax, title=f'Histogram of {column}')
        ax.set_xlabel(column)
        ax.set_ylabel('Frequency')

    plt.tight_layout()
    plt.show()


def calculate_point_biserial_correlation(df, col1, col2):
    """
    计算 DataFrame 中两个指定列之间的点双列相关系数和 p 值。

    参数:
    df (DataFrame): 包含数据的 DataFrame。
    col1 (str): 第一个变量的列名，将被转换为二元变量。
    col2 (str): 第二个变量的列名（连续数值）。

    返回:
    point_biserial_corr (float): 点双列相关系数。
    p_value (float): p 值。
    """
    # 提取列
    col1_data = df[col1]
    col2_data = df[col2]

    # 将col1转换为二元变量
    col1_binary = np.where(col1_data > 0, 1, -1)

    # 计算点双列相关系数
    point_biserial_corr, p_value = pointbiserialr(col1_binary, col2_data)

    # 打印结果
    print(f"{col1} (binary) 和 {col2} 之间的点双列相关系数: {point_biserial_corr}")
    print(f"p值: {p_value}")

    # 可视化
    plt.figure(figsize=(10, 6))
    plt.scatter(col1_binary, col2_data, alpha=0.5)
    plt.xlabel(f'{col1} (binary)')
    plt.ylabel(col2)
    plt.title(f'{col1} (binary) vs {col2}')
    plt.show()

    return point_biserial_corr, p_value